#include "elf.h"
#include "panic.h"
#include "page.h"
#include "pmm.h"
#include "string.h"
#include "proc.h"// cur_proc
#include "list.h"

/* @brief: test magic numbers of elf header
 * @param magics:elf_header->e_ident
 */
void magic_test(unsigned char *magics) {
	if (magics[EI_MAG0] != ELFMAG0
		|| magics[EI_MAG1] != ELFMAG1
		|| magics[EI_MAG2] != ELFMAG2
		|| magics[EI_MAG3] != ELFMAG3
		|| magics[EI_CLASS] != ELFCLASS64
		|| magics[EI_DATA] != ELFDATA2LSB
		|| magics[EI_VERSION] != EV_CURRENT
		|| magics[EI_OSABI] != ELFOSABI_SYSV
		|| magics[EI_ABIVERSION] != 0)
	{
		panic("load_elf: wrong magic number.");
	}
}

/* @brief: transform segment flags to pte flags
 * @param seg_flags: program segment flags
 * @return: pte flags
 */
uint32 flags_trans(uint32 seg_flags) {
	uint32 pte_flags = 0;
	if (seg_flags & PF_X) {
		pte_flags |= PTE_X;
	}
	if (seg_flags & PF_W) {
		pte_flags |= PTE_W;
	}
	if (seg_flags & PF_R) {
		pte_flags |= PTE_R;
	}
	return pte_flags;
}

/* @brief: load all program segments of elf file
 * @param elf_header: the pointer to elf header
 * @param table_ptr3: a level 3 page table that will load the segments
 * @return: program entry(virtual address)
 */
uint64 load_elf(Elf64_Ehdr *elf_header, struct pgtable_l3 *table_ptr3, struct vma **vma_dptr) {
	uint64 file_start = (uint64)elf_header;
	Elf64_Phdr *pro_header;
	int i;
	uint64 size = 0;
	uint64 loaded = 0;
	uint64 offset = 0;
	uint32 pte_flags = 0;
	// physical address
	uint64 pa;
	// virtual address
	uint64 va;
	uint64 pro_addr;
	struct vma *vma_ptr;

	// format test
	magic_test(elf_header->e_ident);
	if (elf_header->e_type != ET_EXEC) {
		panic("load_elf: wrong type.");
	}
	if (elf_header->e_machine != EM_RISCV) {
		panic("load_elf: wrong machine.");
	}
	if (elf_header->e_version != EV_CURRENT) {
		panic("load_elf: wrong version.");
	}
	// load segments
	pro_header = (Elf64_Phdr *)(file_start + (uint64)(elf_header->e_phoff));
	for (i = 0; i < elf_header->e_phnum; ++i, pro_header++) {
		// only segment with LOAD flag can be loaded
		if (pro_header->p_type != PT_LOAD) {
			continue;
		}
		// initialize
		size = 0;
		loaded = 0;
		va = pro_header->p_vaddr;
		// one program segment has the same pte_flags for every page
		pte_flags = flags_trans(pro_header->p_flags) | PTE_U;
		// fill vma
		vma_ptr = (struct vma *)kalloc(sizeof(struct vma));
		vma_ptr->flags = pte_flags;
		vma_ptr->start_addr = va;
		vma_ptr->size = pro_header->p_memsz;
		vma_ptr->next = NULL;
		vma_ptr->count = 1;
		list_add_back((struct linked_list **)vma_dptr, (struct linked_list *)vma_ptr);
		// virtual addrss of current program segment
		pro_addr = file_start + (uint64)(pro_header->p_offset);
		// if va is not PGSIZE aligned, fill 0 for the start of 1st page
		if ((offset = (pro_header->p_vaddr - ROUNDDOWN(pro_header->p_vaddr, PGSIZE))) != 0) {
			pa = kalloc(PGSIZE);
			size = PGSIZE - offset;
			memset((char *)pa, 0, offset);
			memcpy((char *)(pa + offset), (char *)pro_addr, size);
			map_a_page(table_ptr3, ROUNDDOWN(va, PGSIZE), pa, pte_flags);
			loaded += size;
			pro_addr += size;
			va += size;
		}
		// load the rest of pages
		while (loaded < pro_header->p_memsz) {
			pa = kalloc(PGSIZE);
			size = pro_header->p_memsz - loaded > PGSIZE ? PGSIZE : pro_header->p_memsz - loaded;
			memcpy((char *)pa, (char *)pro_addr, size);
			map_a_page(table_ptr3, ROUNDDOWN(va, PGSIZE), pa, pte_flags);
			loaded += size;
			pro_addr += size;
			va += size;
		}
	}
	return elf_header->e_entry;
}
